﻿#include "DriverCallbacks.h"
#include "Defs.h"
#include "Utils.h"
#include <wdm.h>

VOID DriverUnload(__in PDRIVER_OBJECT  driverObject)
{
    DeviceExtension* pdx = static_cast<DeviceExtension*>(driverObject->DeviceObject->DeviceExtension);
    IoDetachDevice(pdx->nextDeviceObject);

    KTIMER kTimer;
    LARGE_INTEGER timeout;
    timeout.QuadPart = 1000000L; //0.1 second
    KeInitializeTimer(&kTimer);
    while(pdx->incompleteIrpsCounter.GetCount() > 0)
    {
        KeSetTimer(&kTimer,timeout,NULL);
        KeWaitForSingleObject(&kTimer,
            Executive,
            KernelMode,
            false,
            NULL);
    }

    IoDeleteDevice(driverObject->DeviceObject);
}

NTSTATUS OnReadCompletion(__in PDEVICE_OBJECT pDeviceObject, __in PIRP pIrp, __in PVOID Context)
{
    DeviceExtension* pdx = static_cast<DeviceExtension*>(pDeviceObject->DeviceExtension);

    if(pIrp->IoStatus.Status == STATUS_SUCCESS)
    {
        const KeyboardInputData* keys = static_cast<KeyboardInputData*>(pIrp->AssociatedIrp.SystemBuffer);
        int numKeys = pIrp->IoStatus.Information / sizeof(KeyboardInputData);

        unsigned int newSize = numKeys;

        //filter keys here
        //utils::FilterKeyboardInputData(keys, sizeof keys, 

        pIrp->IoStatus.Information = newSize * sizeof(KeyboardInputData);

        for(unsigned int i = 0; i < numKeys; ++i)
        {
            if(keys[i].Flags == KEY_MAKE)
            {
                //[TODO]: Get keys list
                //DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL , "KeyDown. ScanCode: %x\n", keys[i].MakeCode);
            }
        }
    }
    if(pIrp->PendingReturned)
    {
        IoMarkIrpPending(pIrp);
    }

    pdx->incompleteIrpsCounter.Dec();
    return pIrp->IoStatus.Status;
}

//Pass any IRP to the next device.
NTSTATUS DispatchThruIO(__in PDEVICE_OBJECT deviceObject, __in PIRP irp)
{
    IoSkipCurrentIrpStackLocation(irp);
    return IoCallDriver((static_cast<DeviceExtension*>(deviceObject->DeviceExtension))->nextDeviceObject, irp);
}

//Pass power management IRP to the next device. For pre-vista OSes only
NTSTATUS DispatchThruPO(__in PDEVICE_OBJECT deviceObject, __in PIRP irp)
{
    ASSERT(!IoIsWdmVersionAvailable(6, 0)); 
    IoSkipCurrentIrpStackLocation(irp);
    return PoCallDriver((static_cast<DeviceExtension*>(deviceObject->DeviceExtension))->nextDeviceObject, irp);
}

NTSTATUS DispatchRead(__in PDEVICE_OBJECT deviceObject, __in PIRP irp)
{
    IoCopyCurrentIrpStackLocationToNext(irp);
    IoSetCompletionRoutine(
        irp,
        OnReadCompletion,
        deviceObject,
        TRUE,
        TRUE,
        TRUE);

    DeviceExtension* pdx = static_cast<DeviceExtension*>(deviceObject->DeviceExtension);
    pdx->incompleteIrpsCounter.Inc();
    return IoCallDriver(pdx->nextDeviceObject, irp);
}
